{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7fc99c3",
   "metadata": {},
   "source": [
    "# Use QR-DQN to Play Pong-v4\n",
    "\n",
    "PyTorch version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2b23c9f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import copy\n",
    "import logging\n",
    "import itertools\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import pandas as pd\n",
    "import gym\n",
    "from gym.wrappers.atari_preprocessing import AtariPreprocessing\n",
    "from gym.wrappers.frame_stack import FrameStack\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "torch.manual_seed(0)\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "        format='%(asctime)s [%(levelname)s] %(message)s',\n",
    "        stream=sys.stdout, datefmt='%H:%M:%S')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "693bdf1e",
   "metadata": {},
   "source": [
    "Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6ab1cb62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:00:46 [INFO] env: <AtariPreprocessing<TimeLimit<AtariEnv<PongNoFrameskip-v4>>>>\n",
      "00:00:46 [INFO] action_space: Discrete(6)\n",
      "00:00:46 [INFO] observation_space: Box(0, 255, (4, 84, 84), uint8)\n",
      "00:00:46 [INFO] reward_range: (-inf, inf)\n",
      "00:00:46 [INFO] metadata: {'render.modes': ['human', 'rgb_array']}\n",
      "00:00:46 [INFO] num_stack: 4\n",
      "00:00:46 [INFO] lz4_compress: False\n",
      "00:00:46 [INFO] frames: deque([], maxlen=4)\n",
      "00:00:46 [INFO] id: PongNoFrameskip-v4\n",
      "00:00:46 [INFO] entry_point: gym.envs.atari:AtariEnv\n",
      "00:00:46 [INFO] reward_threshold: None\n",
      "00:00:46 [INFO] nondeterministic: False\n",
      "00:00:46 [INFO] max_episode_steps: 400000\n",
      "00:00:46 [INFO] _kwargs: {'game': 'pong', 'obs_type': 'image', 'frameskip': 1}\n",
      "00:00:46 [INFO] _env_name: PongNoFrameskip\n"
     ]
    }
   ],
   "source": [
    "env = FrameStack(AtariPreprocessing(gym.make('PongNoFrameskip-v4')),\n",
    "        num_stack=4)\n",
    "for key in vars(env):\n",
    "    logging.info('%s: %s', key, vars(env)[key])\n",
    "for key in vars(env.spec):\n",
    "    logging.info('%s: %s', key, vars(env.spec)[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3747e919",
   "metadata": {},
   "source": [
    "Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5eb74be3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQNReplayer:\n",
    "    def __init__(self, capacity):\n",
    "        self.memory = pd.DataFrame(index=range(capacity),\n",
    "                columns=['state', 'action', 'reward', 'next_state', 'done'])\n",
    "        self.i = 0\n",
    "        self.count = 0\n",
    "        self.capacity = capacity\n",
    "\n",
    "    def store(self, *args):\n",
    "        self.memory.loc[self.i] = args\n",
    "        self.i = (self.i + 1) % self.capacity\n",
    "        self.count = min(self.count + 1, self.capacity)\n",
    "\n",
    "    def sample(self, size):\n",
    "        indices = np.random.choice(self.count, size=size)\n",
    "        return (np.stack(self.memory.loc[indices, field]) for field in\n",
    "                self.memory.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a2b9a656",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, env):\n",
    "        self.action_n = env.action_space.n\n",
    "        self.gamma = 0.99\n",
    "        self.epsilon = 1.\n",
    "\n",
    "        self.replayer = DQNReplayer(capacity=100000)\n",
    "\n",
    "        self.quantile_count = 200\n",
    "        self.cumprob_tensor = torch.arange(1 / (2 * self.quantile_count),\n",
    "                1, 1 / self.quantile_count).view(1, -1, 1)\n",
    "\n",
    "        self.evaluate_net = nn.Sequential(\n",
    "                nn.Conv2d(4, 32, kernel_size=8, stride=4), nn.ReLU(),\n",
    "                nn.Conv2d(32, 64, kernel_size=4, stride=2), nn.ReLU(),\n",
    "                nn.Conv2d(64, 64, kernel_size=3, stride=1), nn.ReLU(),\n",
    "                nn.Flatten(),\n",
    "                nn.Linear(in_features=3136, out_features=512), nn.ReLU(),\n",
    "                nn.Linear(in_features=512,\n",
    "                out_features=self.action_n * self.quantile_count))\n",
    "        self.target_net = copy.deepcopy(self.evaluate_net)\n",
    "        self.optimizer = optim.Adam(self.evaluate_net.parameters(), lr=0.0001)\n",
    "\n",
    "        self.loss = nn.SmoothL1Loss(reduction=\"none\")\n",
    "\n",
    "    def reset(self, mode=None):\n",
    "        self.mode = mode\n",
    "        if mode == 'train':\n",
    "            self.trajectory = []\n",
    "\n",
    "    def step(self, observation, reward, done):\n",
    "        state_tensor = torch.as_tensor(observation,\n",
    "                dtype=torch.float).unsqueeze(0)\n",
    "        q_component_tensor = self.evaluate_net(state_tensor).view(-1,\n",
    "                self.action_n, self.quantile_count)\n",
    "        q_tensor = q_component_tensor.mean(2)\n",
    "        action_tensor = q_tensor.argmax(dim=1)\n",
    "        actions = action_tensor.detach().numpy()\n",
    "        action = actions[0]\n",
    "        if self.mode == 'train':\n",
    "            if np.random.rand() < self.epsilon:\n",
    "                action = np.random.randint(0, self.action_n)\n",
    "            \n",
    "            self.trajectory += [observation, reward, done, action]\n",
    "            if len(self.trajectory) >= 8:\n",
    "                state, _, _, act, next_state, reward, done, _ = \\\n",
    "                        self.trajectory[-8:]\n",
    "                self.replayer.store(state, act, reward, next_state, done)\n",
    "            if self.replayer.count >= 1024 and self.replayer.count % 10 == 0:\n",
    "                self.learn()\n",
    "        return action\n",
    "\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "    def update_net(self, target_net, evaluate_net, learning_rate=0.005):\n",
    "        for target_param, evaluate_param in zip(\n",
    "                target_net.parameters(), evaluate_net.parameters()):\n",
    "            target_param.data.copy_(learning_rate * evaluate_param.data\n",
    "                    + (1 - learning_rate) * target_param.data)\n",
    "\n",
    "    def learn(self):\n",
    "        # replay\n",
    "        batch_size = 32\n",
    "        states, actions, rewards, next_states, dones = \\\n",
    "                self.replayer.sample(batch_size)\n",
    "        state_tensor = torch.as_tensor(states, dtype=torch.float)\n",
    "        reward_tensor = torch.as_tensor(rewards, dtype=torch.float)\n",
    "        done_tensor = torch.as_tensor(dones, dtype=torch.float)\n",
    "        next_state_tensor = torch.as_tensor(next_states, dtype=torch.float)\n",
    "\n",
    "        # compute target\n",
    "        next_q_component_tensor = self.evaluate_net(next_state_tensor).view(\n",
    "                -1, self.action_n, self.quantile_count)\n",
    "        next_q_tensor = next_q_component_tensor.mean(2)\n",
    "        next_action_tensor = next_q_tensor.argmax(dim=1)\n",
    "        next_actions = next_action_tensor.detach().numpy()\n",
    "        all_next_q_quantile_tensor = self.target_net(next_state_tensor\n",
    "                ).view(-1, self.action_n, self.quantile_count)\n",
    "        next_q_quantile_tensor = all_next_q_quantile_tensor[\n",
    "                range(batch_size), next_actions, :]\n",
    "        target_quantile_tensor = reward_tensor.reshape(batch_size, 1) \\\n",
    "                + self.gamma * next_q_quantile_tensor \\\n",
    "                * (1. - done_tensor).reshape(-1, 1)\n",
    "        \n",
    "        all_q_quantile_tensor = self.evaluate_net(state_tensor).view(-1,\n",
    "                self.action_n, self.quantile_count)\n",
    "        q_quantile_tensor = all_q_quantile_tensor[range(batch_size), actions,\n",
    "                :]\n",
    "        \n",
    "        target_quantile_tensor = target_quantile_tensor.unsqueeze(1)\n",
    "        q_quantile_tensor = q_quantile_tensor.unsqueeze(2)\n",
    "        hubor_loss_tensor = self.loss(target_quantile_tensor, q_quantile_tensor)\n",
    "        comparison_tensor = (target_quantile_tensor\n",
    "                < q_quantile_tensor).detach().float()\n",
    "        quantile_regression_tensor = (self.cumprob_tensor\n",
    "                - comparison_tensor).abs()\n",
    "        quantile_huber_loss_tensor = (hubor_loss_tensor\n",
    "                * quantile_regression_tensor).sum(-1).mean(1)\n",
    "        loss_tensor = quantile_huber_loss_tensor.mean()\n",
    "        self.optimizer.zero_grad()\n",
    "        loss_tensor.backward()\n",
    "        self.optimizer.step()\n",
    "\n",
    "        self.update_net(self.target_net, self.evaluate_net)\n",
    "\n",
    "        self.epsilon = max(self.epsilon - 1e-5, 0.05)\n",
    "\n",
    "\n",
    "agent = Agent(env)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f02b76e5",
   "metadata": {},
   "source": [
    "Train & Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34ae8868",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:00:47 [INFO] ==== train ====\n",
      "00:01:04 [DEBUG] train episode 0: reward = -19.00, steps = 1010\n",
      "00:01:34 [DEBUG] train episode 1: reward = -21.00, steps = 852\n",
      "00:02:04 [DEBUG] train episode 2: reward = -21.00, steps = 850\n",
      "00:02:35 [DEBUG] train episode 3: reward = -21.00, steps = 898\n",
      "00:03:11 [DEBUG] train episode 4: reward = -20.00, steps = 989\n",
      "00:03:48 [DEBUG] train episode 5: reward = -20.00, steps = 1040\n",
      "00:04:25 [DEBUG] train episode 6: reward = -20.00, steps = 1035\n",
      "00:05:04 [DEBUG] train episode 7: reward = -19.00, steps = 1058\n",
      "00:05:34 [DEBUG] train episode 8: reward = -21.00, steps = 785\n",
      "00:06:03 [DEBUG] train episode 9: reward = -21.00, steps = 758\n",
      "00:06:41 [DEBUG] train episode 10: reward = -20.00, steps = 990\n",
      "00:07:13 [DEBUG] train episode 11: reward = -20.00, steps = 837\n",
      "00:07:43 [DEBUG] train episode 12: reward = -21.00, steps = 790\n",
      "00:08:18 [DEBUG] train episode 13: reward = -19.00, steps = 964\n",
      "00:08:54 [DEBUG] train episode 14: reward = -19.00, steps = 966\n",
      "00:09:30 [DEBUG] train episode 15: reward = -20.00, steps = 940\n",
      "00:10:13 [DEBUG] train episode 16: reward = -18.00, steps = 1157\n",
      "00:10:48 [DEBUG] train episode 17: reward = -20.00, steps = 921\n",
      "00:11:28 [DEBUG] train episode 18: reward = -18.00, steps = 1048\n",
      "00:12:00 [DEBUG] train episode 19: reward = -21.00, steps = 848\n",
      "00:12:36 [DEBUG] train episode 20: reward = -21.00, steps = 924\n",
      "00:13:19 [DEBUG] train episode 21: reward = -17.00, steps = 1121\n",
      "00:13:55 [DEBUG] train episode 22: reward = -21.00, steps = 910\n",
      "00:14:29 [DEBUG] train episode 23: reward = -20.00, steps = 893\n",
      "00:15:07 [DEBUG] train episode 24: reward = -21.00, steps = 846\n",
      "00:15:42 [DEBUG] train episode 25: reward = -20.00, steps = 870\n",
      "00:16:31 [DEBUG] train episode 26: reward = -17.00, steps = 1248\n",
      "00:17:08 [DEBUG] train episode 27: reward = -21.00, steps = 903\n",
      "00:17:40 [DEBUG] train episode 28: reward = -21.00, steps = 814\n",
      "00:18:17 [DEBUG] train episode 29: reward = -20.00, steps = 899\n",
      "00:18:58 [DEBUG] train episode 30: reward = -20.00, steps = 1004\n",
      "00:19:39 [DEBUG] train episode 31: reward = -20.00, steps = 1052\n",
      "00:20:19 [DEBUG] train episode 32: reward = -20.00, steps = 1023\n",
      "00:20:56 [DEBUG] train episode 33: reward = -21.00, steps = 929\n",
      "00:21:33 [DEBUG] train episode 34: reward = -21.00, steps = 910\n",
      "00:22:09 [DEBUG] train episode 35: reward = -20.00, steps = 942\n",
      "00:22:45 [DEBUG] train episode 36: reward = -21.00, steps = 881\n",
      "00:23:23 [DEBUG] train episode 37: reward = -19.00, steps = 982\n",
      "00:23:52 [DEBUG] train episode 38: reward = -21.00, steps = 762\n",
      "00:24:30 [DEBUG] train episode 39: reward = -21.00, steps = 968\n",
      "00:25:11 [DEBUG] train episode 40: reward = -19.00, steps = 1055\n",
      "00:25:44 [DEBUG] train episode 41: reward = -21.00, steps = 866\n",
      "00:26:15 [DEBUG] train episode 42: reward = -21.00, steps = 799\n",
      "00:26:46 [DEBUG] train episode 43: reward = -21.00, steps = 809\n",
      "00:27:19 [DEBUG] train episode 44: reward = -20.00, steps = 831\n",
      "00:27:53 [DEBUG] train episode 45: reward = -21.00, steps = 882\n",
      "00:28:35 [DEBUG] train episode 46: reward = -19.00, steps = 1095\n",
      "00:29:08 [DEBUG] train episode 47: reward = -20.00, steps = 841\n",
      "00:29:43 [DEBUG] train episode 48: reward = -19.00, steps = 913\n",
      "00:30:16 [DEBUG] train episode 49: reward = -21.00, steps = 849\n",
      "00:30:52 [DEBUG] train episode 50: reward = -20.00, steps = 937\n",
      "00:31:22 [DEBUG] train episode 51: reward = -21.00, steps = 761\n",
      "00:32:03 [DEBUG] train episode 52: reward = -20.00, steps = 1077\n",
      "00:32:39 [DEBUG] train episode 53: reward = -21.00, steps = 910\n",
      "00:33:09 [DEBUG] train episode 54: reward = -21.00, steps = 790\n",
      "00:33:38 [DEBUG] train episode 55: reward = -21.00, steps = 761\n",
      "00:34:14 [DEBUG] train episode 56: reward = -21.00, steps = 942\n",
      "00:34:54 [DEBUG] train episode 57: reward = -19.00, steps = 1032\n",
      "00:35:31 [DEBUG] train episode 58: reward = -19.00, steps = 960\n",
      "00:36:01 [DEBUG] train episode 59: reward = -21.00, steps = 779\n",
      "00:36:39 [DEBUG] train episode 60: reward = -19.00, steps = 1006\n",
      "00:37:16 [DEBUG] train episode 61: reward = -20.00, steps = 962\n",
      "00:37:50 [DEBUG] train episode 62: reward = -21.00, steps = 897\n",
      "00:38:26 [DEBUG] train episode 63: reward = -21.00, steps = 925\n",
      "00:38:57 [DEBUG] train episode 64: reward = -21.00, steps = 806\n",
      "00:39:29 [DEBUG] train episode 65: reward = -21.00, steps = 849\n",
      "00:39:58 [DEBUG] train episode 66: reward = -21.00, steps = 759\n",
      "00:40:28 [DEBUG] train episode 67: reward = -21.00, steps = 787\n",
      "00:41:04 [DEBUG] train episode 68: reward = -21.00, steps = 944\n",
      "00:41:44 [DEBUG] train episode 69: reward = -20.00, steps = 1070\n",
      "00:42:13 [DEBUG] train episode 70: reward = -21.00, steps = 763\n",
      "00:42:42 [DEBUG] train episode 71: reward = -21.00, steps = 762\n",
      "00:43:12 [DEBUG] train episode 72: reward = -21.00, steps = 821\n",
      "00:43:45 [DEBUG] train episode 73: reward = -20.00, steps = 910\n",
      "00:44:20 [DEBUG] train episode 74: reward = -20.00, steps = 971\n",
      "00:44:52 [DEBUG] train episode 75: reward = -21.00, steps = 898\n",
      "00:45:26 [DEBUG] train episode 76: reward = -20.00, steps = 954\n",
      "00:45:59 [DEBUG] train episode 77: reward = -20.00, steps = 917\n",
      "00:46:35 [DEBUG] train episode 78: reward = -20.00, steps = 1024\n",
      "00:47:04 [DEBUG] train episode 79: reward = -21.00, steps = 821\n",
      "00:47:38 [DEBUG] train episode 80: reward = -21.00, steps = 940\n",
      "00:48:10 [DEBUG] train episode 81: reward = -20.00, steps = 898\n",
      "00:48:37 [DEBUG] train episode 82: reward = -21.00, steps = 760\n",
      "00:49:18 [DEBUG] train episode 83: reward = -19.00, steps = 1141\n",
      "00:49:52 [DEBUG] train episode 84: reward = -20.00, steps = 957\n",
      "00:50:33 [DEBUG] train episode 85: reward = -19.00, steps = 1171\n",
      "00:51:12 [DEBUG] train episode 86: reward = -18.00, steps = 1084\n",
      "00:51:41 [DEBUG] train episode 87: reward = -21.00, steps = 834\n",
      "00:52:18 [DEBUG] train episode 88: reward = -21.00, steps = 1020\n",
      "00:52:55 [DEBUG] train episode 89: reward = -20.00, steps = 1035\n",
      "00:53:28 [DEBUG] train episode 90: reward = -21.00, steps = 942\n",
      "00:53:59 [DEBUG] train episode 91: reward = -21.00, steps = 869\n",
      "00:54:30 [DEBUG] train episode 92: reward = -21.00, steps = 869\n",
      "00:55:15 [DEBUG] train episode 93: reward = -18.00, steps = 1261\n",
      "00:55:46 [DEBUG] train episode 94: reward = -21.00, steps = 874\n",
      "00:56:20 [DEBUG] train episode 95: reward = -20.00, steps = 961\n",
      "00:56:58 [DEBUG] train episode 96: reward = -20.00, steps = 1080\n",
      "00:57:32 [DEBUG] train episode 97: reward = -20.00, steps = 945\n",
      "00:58:01 [DEBUG] train episode 98: reward = -21.00, steps = 813\n",
      "00:58:37 [DEBUG] train episode 99: reward = -20.00, steps = 1020\n",
      "00:59:08 [DEBUG] train episode 100: reward = -20.00, steps = 888\n",
      "00:59:48 [DEBUG] train episode 101: reward = -18.00, steps = 1143\n",
      "01:00:27 [DEBUG] train episode 102: reward = -18.00, steps = 1099\n",
      "01:01:06 [DEBUG] train episode 103: reward = -20.00, steps = 1112\n",
      "01:01:41 [DEBUG] train episode 104: reward = -21.00, steps = 985\n",
      "01:02:09 [DEBUG] train episode 105: reward = -21.00, steps = 788\n",
      "01:02:48 [DEBUG] train episode 106: reward = -19.00, steps = 1108\n",
      "01:06:11 [DEBUG] train episode 107: reward = -19.00, steps = 1067\n",
      "01:08:55 [DEBUG] train episode 108: reward = -21.00, steps = 821\n",
      "01:11:30 [DEBUG] train episode 109: reward = -21.00, steps = 782\n",
      "01:14:24 [DEBUG] train episode 110: reward = -21.00, steps = 865\n",
      "01:17:32 [DEBUG] train episode 111: reward = -21.00, steps = 914\n",
      "01:20:54 [DEBUG] train episode 112: reward = -20.00, steps = 974\n",
      "01:24:13 [DEBUG] train episode 113: reward = -20.00, steps = 952\n",
      "01:27:40 [DEBUG] train episode 114: reward = -19.00, steps = 1004\n",
      "01:31:01 [DEBUG] train episode 115: reward = -21.00, steps = 947\n",
      "01:34:35 [DEBUG] train episode 116: reward = -20.00, steps = 1003\n",
      "01:38:20 [DEBUG] train episode 117: reward = -20.00, steps = 1048\n",
      "01:43:38 [DEBUG] train episode 118: reward = -17.00, steps = 1468\n",
      "01:47:22 [DEBUG] train episode 119: reward = -20.00, steps = 1015\n",
      "01:51:21 [DEBUG] train episode 120: reward = -19.00, steps = 1090\n",
      "01:56:15 [DEBUG] train episode 121: reward = -19.00, steps = 1334\n",
      "01:59:55 [DEBUG] train episode 122: reward = -21.00, steps = 986\n",
      "02:05:32 [DEBUG] train episode 123: reward = -17.00, steps = 1550\n",
      "02:09:46 [DEBUG] train episode 124: reward = -18.00, steps = 1204\n",
      "02:14:52 [DEBUG] train episode 125: reward = -16.00, steps = 1456\n",
      "02:18:44 [DEBUG] train episode 126: reward = -21.00, steps = 1113\n",
      "02:24:40 [DEBUG] train episode 127: reward = -16.00, steps = 1707\n",
      "02:28:32 [DEBUG] train episode 128: reward = -17.00, steps = 1126\n",
      "02:31:34 [DEBUG] train episode 129: reward = -21.00, steps = 880\n",
      "02:35:10 [DEBUG] train episode 130: reward = -20.00, steps = 1040\n",
      "02:39:03 [DEBUG] train episode 131: reward = -18.00, steps = 1117\n",
      "02:43:24 [DEBUG] train episode 132: reward = -18.00, steps = 1261\n",
      "02:46:32 [DEBUG] train episode 133: reward = -21.00, steps = 909\n",
      "02:50:02 [DEBUG] train episode 134: reward = -20.00, steps = 1020\n",
      "02:53:46 [DEBUG] train episode 135: reward = -19.00, steps = 1085\n",
      "02:58:00 [DEBUG] train episode 136: reward = -19.00, steps = 1224\n",
      "03:02:48 [DEBUG] train episode 137: reward = -17.00, steps = 1392\n",
      "03:07:44 [DEBUG] train episode 138: reward = -20.00, steps = 1426\n",
      "03:11:50 [DEBUG] train episode 139: reward = -20.00, steps = 1184\n",
      "03:17:01 [DEBUG] train episode 140: reward = -15.00, steps = 1506\n",
      "03:20:48 [DEBUG] train episode 141: reward = -20.00, steps = 1082\n",
      "03:25:04 [DEBUG] train episode 142: reward = -18.00, steps = 1235\n",
      "03:30:03 [DEBUG] train episode 143: reward = -17.00, steps = 1439\n",
      "03:35:44 [DEBUG] train episode 144: reward = -17.00, steps = 1645\n",
      "03:41:22 [DEBUG] train episode 145: reward = -19.00, steps = 1626\n",
      "03:46:15 [DEBUG] train episode 146: reward = -17.00, steps = 1409\n",
      "03:51:03 [DEBUG] train episode 147: reward = -18.00, steps = 1383\n",
      "03:57:08 [DEBUG] train episode 148: reward = -17.00, steps = 1744\n",
      "04:01:27 [DEBUG] train episode 149: reward = -21.00, steps = 1238\n",
      "04:05:28 [DEBUG] train episode 150: reward = -20.00, steps = 1142\n",
      "04:11:23 [DEBUG] train episode 151: reward = -15.00, steps = 1681\n",
      "04:16:41 [DEBUG] train episode 152: reward = -19.00, steps = 1495\n",
      "04:22:18 [DEBUG] train episode 153: reward = -15.00, steps = 1572\n",
      "04:27:09 [DEBUG] train episode 154: reward = -17.00, steps = 1366\n",
      "04:32:09 [DEBUG] train episode 155: reward = -18.00, steps = 1398\n",
      "04:37:54 [DEBUG] train episode 156: reward = -16.00, steps = 1607\n",
      "04:45:05 [DEBUG] train episode 157: reward = -11.00, steps = 1989\n",
      "04:50:59 [DEBUG] train episode 158: reward = -16.00, steps = 1626\n",
      "04:57:05 [DEBUG] train episode 159: reward = -15.00, steps = 1664\n",
      "05:03:14 [DEBUG] train episode 160: reward = -15.00, steps = 1608\n",
      "05:09:09 [DEBUG] train episode 161: reward = -18.00, steps = 1579\n",
      "05:14:15 [DEBUG] train episode 162: reward = -18.00, steps = 1354\n",
      "05:20:52 [DEBUG] train episode 163: reward = -13.00, steps = 1744\n",
      "05:26:50 [DEBUG] train episode 164: reward = -16.00, steps = 1562\n",
      "05:33:57 [DEBUG] train episode 165: reward = -16.00, steps = 1860\n",
      "05:41:52 [DEBUG] train episode 166: reward = -15.00, steps = 2063\n",
      "05:48:29 [DEBUG] train episode 167: reward = -14.00, steps = 1722\n",
      "05:55:25 [DEBUG] train episode 168: reward = -14.00, steps = 1803\n",
      "06:04:29 [DEBUG] train episode 169: reward = -7.00, steps = 2342\n",
      "06:12:12 [DEBUG] train episode 170: reward = -9.00, steps = 1988\n",
      "06:21:08 [DEBUG] train episode 171: reward = -8.00, steps = 2270\n",
      "06:29:53 [DEBUG] train episode 172: reward = -6.00, steps = 2237\n",
      "06:38:32 [DEBUG] train episode 173: reward = -8.00, steps = 2209\n",
      "06:47:39 [DEBUG] train episode 174: reward = -7.00, steps = 2319\n",
      "06:52:05 [DEBUG] train episode 175: reward = -19.00, steps = 1122\n",
      "06:57:46 [DEBUG] train episode 176: reward = -14.00, steps = 1435\n",
      "07:04:08 [DEBUG] train episode 177: reward = -12.00, steps = 1604\n",
      "07:12:55 [DEBUG] train episode 178: reward = -8.00, steps = 2213\n",
      "07:19:41 [DEBUG] train episode 179: reward = -13.00, steps = 1704\n",
      "07:26:31 [DEBUG] train episode 180: reward = -11.00, steps = 1704\n",
      "07:31:32 [DEBUG] train episode 181: reward = -16.00, steps = 1256\n",
      "07:35:27 [DEBUG] train episode 182: reward = -19.00, steps = 983\n",
      "07:39:08 [DEBUG] train episode 183: reward = -20.00, steps = 919\n",
      "07:42:53 [DEBUG] train episode 184: reward = -19.00, steps = 938\n",
      "07:50:35 [DEBUG] train episode 185: reward = -9.00, steps = 1921\n",
      "07:58:47 [DEBUG] train episode 186: reward = -8.00, steps = 2045\n",
      "08:03:41 [DEBUG] train episode 187: reward = -16.00, steps = 1219\n",
      "08:08:19 [DEBUG] train episode 188: reward = -18.00, steps = 1151\n",
      "08:16:34 [DEBUG] train episode 189: reward = -9.00, steps = 2047\n",
      "08:25:00 [DEBUG] train episode 190: reward = -7.00, steps = 2084\n",
      "08:34:15 [DEBUG] train episode 191: reward = -5.00, steps = 2298\n"
     ]
    }
   ],
   "source": [
    "def play_episode(env, agent, max_episode_steps=None, mode=None, render=False):\n",
    "    observation, reward, done = env.reset(), 0., False\n",
    "    agent.reset(mode=mode)\n",
    "    episode_reward, elapsed_steps = 0., 0\n",
    "    while True:\n",
    "        action = agent.step(observation, reward, done)\n",
    "        if render:\n",
    "            env.render()\n",
    "        if done:\n",
    "            break\n",
    "        observation, reward, done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        elapsed_steps += 1\n",
    "        if max_episode_steps and elapsed_steps >= max_episode_steps:\n",
    "            break\n",
    "    agent.close()\n",
    "    return episode_reward, elapsed_steps\n",
    "\n",
    "\n",
    "logging.info('==== train ====')\n",
    "episode_rewards = []\n",
    "for episode in itertools.count():\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent, mode='train')\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('train episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "    if np.mean(episode_rewards[-5:]) > 16.:\n",
    "        break\n",
    "plt.plot(episode_rewards)\n",
    "\n",
    "\n",
    "logging.info('==== test ====')\n",
    "episode_rewards = []\n",
    "for episode in range(100):\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('test episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "logging.info('average episode reward = %.2f ± %.2f',\n",
    "        np.mean(episode_rewards), np.std(episode_rewards))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
